MemGPT 源码阅读
 · 阅读需 29 分钟
1. 整体架构概览
1.1 MemGPT的核心模块及其职责
graph LR
    A[agent] --> B[agent_store]
    A --> C[data_types]
    A --> D[embeddings]
    A --> E[functions]
    A --> F[memory]
    A --> G[models]
    A --> H[persistence_manager] 
    A --> I[prompts]
    A --> J[server]
agent: Agent的核心逻辑,包括初始化、处理消息、状态管理等agent_store: Agent状态的持久化存储,支持不同的后端(如PostgreSQL、Chroma等)data_types: MemGPT使用的核心数据结构,如AgentState、Message、Passage等embeddings: 封装embedding相关逻辑,如文本分块、向量化等functions: 自定义函数的定义、注册与调用memory: 包括CoreMemory、RecallMemory和ArchivalMemory,分别对应不同层次的记忆models: 定义了与LLM交互的请求与响应格式persistence_manager: 管理Agent状态在不同存储后端间的持久化prompts: 包含了一些默认的prompt模板server: 提供了将Agent封装为RESTful API服务的能力
1.2 模块交互与数据流
sequenceDiagram
    participant C as Client
    participant S as Server
    participant A as Agent
    participant M as Memory
    participant P as PersistenceManager
    participant L as LLM
    
    C->>S: 发送消息
    S->>A: 将消息传递给Agent
    A->>M: 查询/更新 Memory
    A->>P: 持久化状态变更
    A->>L: 发送消息序列,请求LLM响应
    L-->>A: 返回LLM响应
    A->>A: 解析LLM响应(含函数调用)
    A-->>S: 返回最终回复
    S-->>C: 返回最终回复
- Client通过Server将消息发送给Agent
 - Agent查询/更新内存,并通过PersistenceManager持久化状态变更
 - Agent将消息历史发送给LLM,请求下一步响应
 - LLM返回响应,其中可能包含函数调用
 - Agent解析LLM响应,必要时进行函数调用,最终生成回复
 - Agent通过Server将最终回复返回给Client
 
2. Agent的状态管理
2.1 Agent类的设计与实现
Agent类(agent.py)是MemGPT的核心,它封装了一个Agent的完整状态和交互逻辑。其主要组成部分包括:
config: Agent的配置,如LLM/Embedding设置、所使用的memory/function等memory: Agent的核心记忆,包括persona和human部分functions: Agent可调用的函数定义messages: Agent的历史消息记录
Agent类的核心方法是step(),它接受一条用户消息,经过一系列处理步骤,最终生成Agent的回复:
- 将用户消息添加到
messages - 构造prompt(包括
system、memory、functions、messages等),发送给LLM - 解析LLM响应,提取
函数调用和回复内容 - 如果存在
函数调用,则执行对应的Python函数,将函数结果添加到messages - 如有必要,总结
messages以控制token数量 - 将Agent回复添加到
messages - 根据需要持久化状态变更(通过
PersistenceManager) 
2.2 PersistenceManager与存储后端
PersistenceManager抽象类(persistence_manager.py)定义了一组接口,用于管理Agent状态在不同存储后端间的持久化。
目前主要的实现是LocalStateManager,它将状态直接维护在内存中,主要用于demo/debug场景。对于生产环境,则需要实现基于数据库的PersistenceManager,将状态持久化到数据库中。
PersistenceManager的主要方法包括:
append_messages(): 添加新消息trim_messages(): 从头部删除部分消息swap_system_message(): 更换system消息update_memory(): 更新memory内容
通过PersistenceManager,可以将Agent的状态变更同步到底层存储,并在Agent重启后恢复之前的状态。
2.3 AgentState等核心数据结构
AgentState(data_types.py)封装了一个Agent的完整状态,包括其配置、初始prompt、函数定义、消息历史等。AgentState主要用于在Agent和持久化存储之间传递状态。
其他一些核心数据结构还包括:
Message: 表示Agent发送或接收的一条消息,包括role(system/assistant/user)、content、function_call等字段Passage: 表示一段用于语义检索的文本片段,包括text、embedding等字段LLMConfig/EmbeddingConfig: 封装LLM和Embedding的配置,如model、api_key等
2.4 Agent运行过程中的状态变迁
下面以一个典型的场景为例,展示Agent运行过程中的状态变迁:
stateDiagram-v2
    [*] --> Idle: 创建Agent
    Idle --> WaitingForUserMessage: 等待用户消息
    WaitingForUserMessage --> ProcessingUserMessage: 收到用户消息
    ProcessingUserMessage --> WaitingForLLMResponse: 发送消息给LLM
    WaitingForLLMResponse --> ProcessingLLMResponse: 收到LLM响应
    ProcessingLLMResponse --> CallingFunction: 调用函数
    CallingFunction --> ProcessingLLMResponse: 函数执行完毕
    ProcessingLLMResponse --> GeneratingReply: 生成回复
    GeneratingReply --> PersistingState: 持久化状态变更
    PersistingState --> WaitingForUserMessage: 等待下一条用户消息
- Agent创建后处于
Idle状态,等待用户消息 - 收到用户消息后,进入
ProcessingUserMessage状态,将消息添加到messages - 发送消息序列给LLM,进入
WaitingForLLMResponse状态 - 收到LLM响应后,进入
ProcessingLLMResponse状态,解析响应 - 如果响应中包含
函数调用,则进入CallingFunction状态,执行对应函数 - 函数执行完毕后,回到
ProcessingLLMResponse状态 - 生成最终回复,进入
GeneratingReply状态 - 持久化状态变更(
messages、memory等),进入PersistingState状态 - 状态持久化完成后,回到
WaitingForUserMessage状态,等待下一条用户消息